import pandas as pd
import os
import time

class ResultManager:
    def __init__(self):
        self.epislodes_num = None
        self.result = {}
        self.success_cnt = 0
        self.columns = []
        self.result_dir = "resultlib"
        if not os.path.exists(self.result_dir):
            os.makedirs(self.result_dir)
    
    def init_result(self, episodes_num):
        self.episodes_num = episodes_num
        self.success_cnt = 0
        self.result.clear()
        self.columns.clear()

    def result_push_in(self, episode_idx, **kwargs):
        if self.episodes_num is None:
            raise ValueError("ResultManager has not been initialized with episodes_num.")
        
        if self.result.get("episode_idx") is None: # Initialize when result is empty 
            self.result["episode_idx"] = [None] * self.episodes_num
            for key in kwargs.keys():
                self.result[key] = [None] * self.episodes_num
                self.columns.append(key)
        
        for key, value in kwargs.items():
            if key not in self.columns:
                self.columns.append(key)
                self.result[key] = [None] * self.episodes_num
                print("[ResultManager] Warning: New column added:", key)
            
            self.result[key][episode_idx] = value
    
    def save_result(self, dir):
        filename = f"{dir}/test_result.csv"
        
        df = pd.DataFrame(self.result, columns=self.columns)
        df.to_csv(filename)

        self.episodes_num = None
    def clear_result(self):
        self.episodes_num = None
        self.success_cnt = 0
        self.result.clear()
        self.columns.clear()